#!/usr/bin/env python
'''
   load.py  -  Loads a dictionary lookup-file into memory, ready to aid 
               in translation lookups

   Tim Kane (385199)
   Baolin Jing (2719650)
   Zhe Zhang (304914)
 
   19 March 2010
'''

import sys
import os
import re
#import pprint
#import nltk

from AlignedSent import AlignedSent
from demo import demo


def load(filename):
    '''
    loads the set of files identified by 'filename'
    that is, it will load trial.e, trial.f, wrial.wa


    sentence_no position_L1 position_L2 [S|P] [confidence]

    where:
        - sentence_no represents the id of the sentence within the test file.
    Sentences in the test data already have an id assigned. (see the examples
    below)
    
        - position_L1 represents the position of the token that is aligned from
    the text in language L1; the first token in each sentence is token 1. (not 0)
    
        - position_L2 represents the position of the token that is aligned from the
    text in language L2; again, the first token is token 1.
    
        - S|P can be either S or P, representing a Sure or Probable alignment. All
    alignments that are tagged as S are also considered to be part of the P
    alignments set (that is, all alignments that are considered "Sure" alignments
    are also part of the "Probable" alignments set). If the S|P field is missing,
    a value of S will be assumed by default.
    
        - confidence is a real number, in the range (0-1] (1 meaning highly confident,
    0 meaning not confident); this field is optional, and by default we assume a
    confidence number of 1.
    
    
    Running example

    ---------------
    
    Alternatively, systems may also provide an S|P marker and/or a confidence
    score, as in the following example:
    
    18 1 1 1
    18 2 2 P 0.7
    18 3 3 S
    18 4 4     S 1
        
    with missing S|P fields considered by default to be S, and missing confidence
    scores considered by default 1.

        @type  filename: string
        @param filename: the file name (including path) for the language files
        @return:         a dictionary containing AlignedSent objects
    
    '''
    
    translation = dict()


    # This isn't the best way to load the translation data.

    # This loop structure will initially load ALL files into memory, indexed 
    # on ID.  Then for each ID, create the AlignedSent object in one hit.  
    # This is easier than trying to create objects and then 'update' them as we
    # find more data.
    # This will be done at the expense of memory to retain all transation files 
    # in mem at startup for a short period

    
    mem = dict()
    for ext in ['e', 'f', 'wa']:
        mem[ext] = dict()

        fname = filename + '.' + ext

        if (os.path.isfile(fname)):
            infile = open(fname,'r')
        else:
            print ("invalid path/file: %s" % fname)
            exit(-1)

        print "importing from extension " + ext

        if (ext == 'wa'):
            # It's a word alignment file

            # This is the format we expect (seperated by a single whitespace)
            #
            #    S   L1  L2  S|P Confidence
            #    18  1   1   1
            #    18  2   2   P   0.7
            #    18  3   3   S
            #    18  4   4   S   1

            for line in infile:
        
                pattern = r'^(\d+)\W(\d+)\W(\d+)(?:\W([S|P]))?\W([0-9\.]+)?'
                search = re.findall (pattern, line)

                index = int(search[0][0])

                if index in mem[ext]:
                    # Key exists - extend
                    mem[ext][index].append (search[0][1:])
                else:
                    # This is a new key
                    item = [ search[0][1:] ] 
                    mem[ext][index] = item
        else:
            # This is the format we expect (seperated by a single whitespace)
            #
            #    <s snum=1001> each of them is very complex , the end .  </s>

            for line in infile:

                pattern = r'^<s snum=([0-9]+)>(.*)</s>$'
                search = re.findall (pattern, line)

                index  = int(search[0][0])
                string = search[0][1].strip()

                mem[ext][index] = string



        infile.close()


    # So we've got everything in memory - let's instantiate it all correctly
    # Assume that every sentence defined in t1 and t2 also has a word-alignment 
    # defined.
    #
    # Use the list of sentence id's found in the word-alignments to iterate 
    # through each set of sentences

    print "aligning sentences"


    for index in mem['wa']:
        #print mem['e'][index]
        #print mem['f'][index]
        #print mem['wa'][index]

            # Ready to initialise a new instance with this data
            # (or update an existing instance?)

            # In this case, we only care about 'SURE' translations.
        
        sentence = AlignedSent (mem['e'][index], 
                                mem['f'][index], 
                                mem['wa'][index], 'S')

        translation[index] = sentence

    return translation


if __name__ == "__main__":
    '''
        main function - demonstrate the functionality of the AlignedSent class
    '''
    #pylint: disable-msg=C0103

    args = sys.argv
    if len(args) != 2:
        print 'path name should be specified after the program name'
        exit(0)
    for arg in args:
        if 'load.py' in arg:
            continue
        else:
            filname = arg

    translation = load(filname)
    demo(translation)

